import numpy as np
import matplotlib.pyplot as plt
data=[[3, 104, 0],
     [2, 100, 0],
     [1, 81, 0],
     [101, 10, 1],
     [99, 5, 1],
     [98, 2, 1]]
#1.数据处理
data=np.array(data)
x_test=[58,20]   #测试数据
x_test=np.array(x_test)
plt.scatter(x_test[0],x_test[1])
plt.scatter(data[:,0],data[:,1],c=data[:,-1].ravel())
plt.show()

def knn_w(data,x_test,k):
    dist=[]
    for sample in data:
        sampledist=np.sqrt(np.sum((sample[:-1]-x_test)**2))
        dist.append([sampledist,sample[-1]])
    dist.sort()
    dist_k=np.array(dist[:k])
    label_k=dist_k[:,-1]  #只提取出前k个归属
    dist_s=dist_k[:,:-1]  #只提取出前k个距离
    print(label_k)
    w0=0   #初始0类的权重
    w1=0   #初始1类的权重
    for label,dd in zip(label_k,dist_s):   #序列解包，双重遍历方式
        if label==1:
            w1+=1/dd   #权重选取方式为距离的倒数
        else:
            w0+=1/dd
    if w1>w0:   #比较权重w0,w1,的大小，谁权重大归哪类
        print('预测',1)
    else:
        print('预测',0)

knn_w(data,x_test,5)

